#!/usr/bin/python
import sys, os, re, time
################################################################################
################################################################################
class Params:
   verbose = True
   rebuild = False
   mainmem = 96*1024*1024*1024
   memfact = 40
   threads = [ 24, 16, 8, 4, 2, 1 ]
   def __init__( self, mtxpath ):
      
      assert os.path.isdir( mtxpath )
      self.mtxpath = mtxpath
      
      self.formats = [ 'RMA', 'QTS', 'NZD', 'CSR', 'COO' ]
      self.degrees = [ 2, 4, 8 ]
      
      dagHomeEnvName = 'DAGMTX_HOME_DIR'
      
      if not dagHomeEnvName in os.environ:
         print 'Before using ' + sys.argv[0] + ', you need to set the environment variable "' + dagHomeEnvName + '"'
         print 'to point to the home directory for this project.'
         sys.exit( -1 )
      
      self.dagHomeDir = os.environ[ dagHomeEnvName ]      
      self.pysHome = os.path.join( self.dagHomeDir, 'pys' )
      self.buildPy = os.path.join( self.pysHome, 'BuildDAGMtx.py' )
   

################################################################################
################################################################################
def AddRedirectionToCmd( cmd, fn ):
  return " ".join([ cmd, "1>", fn, "2>&1" ])

def PrintTimeStamp():
   theTime = time.localtime()
   yr  = '%04d' % theTime.tm_year
   mon = '%02d' % theTime.tm_mon
   day = '%02d' % theTime.tm_mday
   hr  = '%02d' % theTime.tm_hour
   mt  = '%02d' % theTime.tm_min
   sec = '%02d' % theTime.tm_sec
   print '...', '/'.join( [yr, mon, day] ), ':'.join( [hr, mt, sec] )
   sys.stdout.flush()

def SystemWrapperPipelined( cmd, verbose, trial = False ):
   if verbose:
      PrintTimeStamp()
      print '... os.popen(', cmd, ')'
      sys.stdout.flush()
   if not trial:
      FH = os.popen( cmd )
      output = [ line.rstrip() for line in FH.readlines() ]
      FH.close()
      return output
   return []

def SystemWrapper( cmd, verbose, trial = False ):
   dt = 0
   if verbose:
      PrintTimeStamp()
      print '... os.system(', cmd, ')'
      sys.stdout.flush()
   if trial:
      status = 0
   else:
      t0 = time.time()
      status = os.system( cmd )
      t1 = time.time()
      dt = t1 - t0
   assert status == 0
   if verbose:
      print '...', dt, 'seconds elapsed during system call.'
   return dt

def RemoveWrapper( fn, verbose, trial = False ):
	if verbose:
		PrintTimeStamp()
		print '... os.remove(', fn, ')'
		sys.stdout.flush()
	if not trial:
		os.remove( fn )

################################################################################
################################################################################
def HeadNLines( fn, n ):
   cmd = ' '.join([ 'head', '-%d' % n, fn ])
   return SystemWrapperPipelined( cmd, False, False )

def IsRealMtx( mtxFile ):
	line = HeadNLines( mtxFile, 1 )[0].rstrip().lstrip( '%' )
	toks = line.split()
	if len( toks ) != 5:          return False
	if toks[0] != 'MatrixMarket': return False
	if toks[1] != 'matrix':       return False
	if toks[2] != 'coordinate':   return False
	if toks[3] != 'real':         return False
	if toks[4] not in [ 'general', 'symmetric' ]: return False
	return True

def MtxDims( mtxFile ):
   lines = HeadNLines( mtxFile, 200 )
   for line in lines:
      if len( line ):
         if line[0] != '%': break
   toks = line.split()
   assert len( toks ) == 3
   m   = int( toks[0] )
   n   = int( toks[1] )
   nnz = int( toks[2] )
   return ( m, n, nnz )

def CSRBytes( m, n, nnz ):
   nnzValBytes = 8 * nnz
   colIdxBytes = 4 * nnz
   rowPtrBytes = 4 * m + 4
   totalBytes = nnzValBytes + colIdxBytes + rowPtrBytes
   return totalBytes

################################################################################
################################################################################
def ForkLaunch( cmds, logs, maxFork = 1, verbose = True, trial = False ):
   # sanity:
   assert len( cmds ) == len( logs )
   #
   # setup:
   cmdslogs = zip( cmds, logs )
   numLaunched = 0
   #
   # fork & launch:
   for ( cmd, log ) in cmdslogs:
      pid = os.fork()
      cmd = AddRedirectionToCmd( cmd, log )
      if pid == 0:
         SystemWrapper( cmd, verbose = verbose, trial = trial )
         sys.exit( 0 )
      else:
         numLaunched += 1
         if numLaunched >= maxFork:
            os.waitpid( -1, 0 )
   while True:
      try:    os.waitpid( -1, 0 )
      except: break

################################################################################
################################################################################
def DagFileExists( mtxFile, fmt, deg ):
   parts = 1
   ( path, fn ) = os.path.split( mtxFile )
   toks = fn.split( '.' )
   toks.pop()
   rawf = '.'.join( toks )
   rawdag = '%s.%s.deg%02d.parts%02d.dag' % ( rawf, fmt.lower(), deg, parts )
   dagfn  = os.path.join( path, rawdag )
   if os.path.isfile( dagfn ): return True
   return False

def MemLimit( params, numThreads ):
   if numThreads == 1: return 1024*1024*1024*1024
   return params.mainmem / numThreads / params.memfact

def GenCmdLogTup( params, mtxFile, fmt, deg ):
   parts = 1
   ( path, fn ) = os.path.split( mtxFile )
   toks  = fn.split( '.' )
   rawf  = toks[0]
   ( path, fn ) = os.path.split( params.buildPy )
   toks  = fn.split( '.' )
   rawc  = toks[0]
   cmd   = ' '.join( [ params.buildPy, fmt, mtxFile, str( deg ), str( parts ) ] )
   log   = '.'.join( [           rawc, fmt,    rawf, str( deg ), str( parts ), 'log' ] )
   log   = os.path.join( params.mtxpath, log )
   return ( cmd, log )

def GenerateCmdAndLogList( params, mtxFileTups, numThreads ):
   cmdLogList = []
   memLimit = MemLimit( params, numThreads )
   while len( mtxFileTups ):
      tup = mtxFileTups.pop()
      ( csrBytes, mtxFile, m, n, nnz ) = tup
      if csrBytes < memLimit:
         for fmt in params.formats:
            for deg in params.degrees:
               if params.rebuild or not DagFileExists( mtxFile, fmt, deg ):
                  cmdLogTup = GenCmdLogTup( params, mtxFile, fmt, deg )
                  cmdLogList.append( cmdLogTup )
      else:
         mtxFileTups.append( tup )
         break
   return cmdLogList

def Convert2DAGs( params, mtxFileTups ):
   for numThreads in params.threads:
      cmdslogs = GenerateCmdAndLogList( params, mtxFileTups, numThreads )
      if cmdslogs == []: continue
      ( cmds, logs ) = zip( *cmdslogs )
      ForkLaunch( cmds, logs, maxFork = numThreads, verbose = params.verbose, trial = False )

def SortAndSizeMtxs( mtxFiles ):
   tups = []
   for mtxFile in mtxFiles:
      ( m, n, nnz ) = MtxDims( mtxFile )
      csrBytes = CSRBytes( m, n, nnz )
      tup = ( csrBytes, mtxFile, m, n, nnz )
      tups.append( tup )
   tups.sort()
   tups.reverse()
   return tups

def WalkAndFindMtxs( params ):
   mtxFiles = []
   for ( dirpath, dirs, files ) in os.walk( params.mtxpath ):
      for fn in files:
         toks = fn.split( '.' )
         if toks[ -1 ] == 'mtx':
            mtxFile = os.path.join( dirpath, fn )
            if IsRealMtx( mtxFile ):
               mtxFiles.append( mtxFile )
   return mtxFiles

################################################################################
################################################################################
if len( sys.argv ) < 2:
   ( progPath, progName ) = os.path.split( sys.argv[0] )
   print 'Usage:', progName, '<pathToMtxs>'
   sys.exit( -1 )

params = Params( sys.argv[1] )

mtxFiles    = WalkAndFindMtxs( params )
mtxFileTups = SortAndSizeMtxs( mtxFiles )
Convert2DAGs( params, mtxFileTups )
